#ifndef DISJOINTSET_H
#define DISJOINTSET_H

#include <stdlib.h>
#include <memory.h>

/* Disjoint-set forest data structure. */
typedef struct DisjointSet
{
    int *parent;            /* index to parent node; 0 if root */
    unsigned char *rank;    /* depth of children before path compression */
    int num_items;          /* number of items */
    int num_disjoint;       /* number of disjoint subsets */
} DisjointSet;

/* Creates a disjoint set. */
void DisjointSet_create(DisjointSet *set, int num_items)
{
    set->num_items = num_items;
    set->num_disjoint = num_items;
    set->parent = calloc(num_items + 1, sizeof(set->parent[0]));
    set->rank = calloc(num_items + 1, sizeof(set->rank[0]));
}

/* Destroys a disjoint set. */
void DisjointSet_destroy(DisjointSet *set)
{
    free(set->parent);
    free(set->rank);
}

/* Finds the representative of a given element. Both the input index and 
 * the returned index are one-based. If 'item' = 0, returns 0.
 */
int DisjointSet_find(DisjointSet *set, int item)
{
    int parent = set->parent[item];
    if (item == 0)                  /* null item */
        return 0;
    if (parent == 0)                /* item is root */
        return item;
    if (set->parent[parent] == 0)   /* parent is root (path compressed) */
        return parent;
    return set->parent[item] = DisjointSet_find(set, parent);
}

/* Checks whether two items are in the same class. If one of x and y is 
 * zero and the other is non-zero, they are considered not connected. 
 * If both x and y are zero, they are considered not connected.
 */
int DisjointSet_connected(DisjointSet *set, int x, int y)
{
    if (x == 0 || y == 0)
        return 0;
    else
        return DisjointSet_find(set, x) == DisjointSet_find(set, y);
}

/* Merges two subsets. Returns the representative of the merged sets. 
 * If either is zero, returns the other. If both are zero, returns 0.
 */
int DisjointSet_merge(DisjointSet *set, int x, int y)
{
    x = DisjointSet_find(set, x);
    y = DisjointSet_find(set, y);
    if (x == y)
        return x;
    if (x == 0)
        return y;
    if (y == 0)
        return x;
    
    --set->num_disjoint;        /* update number of disjoint sets */

    if (set->rank[x] < set->rank[y])
    {
        set->parent[x] = y;     /* append x to y */
        return y;
    }
    else if (set->rank[x] > set->rank[y])
    {
        set->parent[y] = x;     /* append y to x */
        return x;
    }
    else
    {
        set->parent[x] = y;     /* append x to y (arbitrary) */
        set->rank[y]++;         /* increase rank */
        return y;
    }
}

#endif /* DISJOINTSET_H */
